#!/usr/bin/python
# 将模型封装服务，并进行自测。
# 2022-12-9
from model_load import *

# cgec 中文语法纠错模块
def cgec_process(question = cgec_question):
    cgec_result = cgec_model(question)
    return cgec_result["output"]

# egec 英文语法纠错模块
def egec_process(question = egec_question):
    result = egec_model.generate_text("grammar: {}".format(question), args=args)
    return result.text

# # eaqg 英文有答案问题生成
# def eaqg_process(text = eaqg_text):
#     result = eaqg_model(text)
#     return result[0]["generated_text"]
#
# # caqg 中文有答案问题生成
# def caqg_process(text=caqg_text):
#     inputs = caqg_tokenizer.encode_plus(text,max_length=448,padding="max_length",truncation=True,return_tensors='pt')
#     output = caqg_model.generate(input_ids=inputs['input_ids'],attention_mask=inputs['attention_mask'],do_sample=True,num_beams=5,max_length=64,top_p = 0.9,)
#     result = caqg_tokenizer.batch_decode(output,clean_up_tokenization_spaces=True, skip_special_tokens=True)[0].replace("问题:","")
#     return result

# cnaqg 中文无答案问题生成
def cnaqg_process(text=cnaqg_text):
    temp = []
    temp.append(text)
    result = cnaqg_model(temp)
    return result[0]["generated_text"].split("<sep>")[0]

# enaqg 英文无答案问题生成
def enaqg_process(text=enaqg_text):
    inputs = enaqg_tokenizer(text, return_tensors="pt")
    output = enaqg_model.generate(**inputs, max_length=40)
    res = enaqg_tokenizer.decode(output[0], skip_special_tokens=True)
    return res

# edg 英文干扰项生成
def edg_process(context=edg_context,
                question=edg_question,
                answer=edg_answer):
    en_dg_text = context + " </s> " + question + " </s> " +answer
    distractor = edg_model(en_dg_text)[0]["generated_text"]
    return distractor

# cdg 中文干扰项生成
def cdg_process(answer=cdg_answer):
    distractor = cdg_model.nearby(answer, 3)[0]
    distractor = str(distractor).replace("[","").replace("]","").replace("'","").replace(",","$")
    return distractor.replace(" ","")

# eqa 英文QA评估可回答性模块
def eqa_process(context=eqa_context,
                question=eqa_question):
    context_question = {"context": context, "question": question}
    answer = eqa_model(context_question)['answer']
    answer_confidence = eqa_model(context_question)['score']
    return answer, answer_confidence

# cqa 英文QA评估可回答性模块
def cqa_process(context=cqa_context,
                question=cqa_question):
    context_question = {"context": context, "question": question}
    answer = cqa_model(context_question)['answer']
    answer_confidence = cqa_model(context_question)['score']
    return answer, answer_confidence

# ceval 中文 一致性相关性评估模块
def ceval_process(context=ch_eval_context,
                  question=ch_eval_question,
                  reference=ch_eval_ref):
    consistency = scorer.score(doc=context, refs=[], hypo=question, aspect='consistency')
    relevance = scorer.score(doc=context, refs=[reference], hypo=question, aspect='relevance')
    if reference == ch_eval_ref:
        if question != ch_eval_question:
            return {"consistency": consistency}
    return {"consistency": consistency,
            "relevance": relevance}

# eeval 英文 一致性相关性评估模块
def eeval_process(context=en_eval_context,
                  question=en_eval_question,
                  reference=en_eval_ref):
    consistency = scorer.score(doc=context, refs=[], hypo=question, aspect='consistency')
    relevance = scorer.score(doc=context, refs=[reference], hypo=question, aspect='relevance')
    if reference == en_eval_ref:
        if question != en_eval_question:
            return {"consistency": consistency}
    return {"consistency": consistency,
            "relevance": relevance}

# # ener 英文命名实体识别模块
def ener_process(text=en_ner_text):
    doc = en_ner_model(text)
    res = []
    for i in doc.entities:
        i = str(i)
        i = i.replace("{","").replace("}","").replace('"',"").replace("text","").replace(":","").replace("\n","").replace(" ","").split(",")[0]
        res.append(i)
    return {"entities": res}

# cner 中文命名实体识别模块
def cner_process(text=ch_ner_text):
    doc = ch_ner_model(text)
    res = []
    for i in doc.entities:
        i = str(i)
        i = i.replace("{","").replace("}","").replace('"',"").replace("text","").replace(":","").replace("\n","").replace(" ","").split(",")[0]
        res.append(i)
    return {"entities": res}

# 模型初始化测试
def model_self_test():
    init_result = {}
    # 测试中文语法纠错模块
    cgec_result = cgec_process()
    if cgec_result == "遇到逆境时，我们必须勇于面对，而且要愈挫愈勇，这样我们才能朝着成功之路前进。":
        init_result["中文语法纠错模块"] = "初始化成功"
    else:
        init_result["中文语法纠错模块"] = "初始化失败"
    print("cgec",init_result["中文语法纠错模块"])
    # 测试英文语法纠错模块
    egec_result = egec_process()
    if egec_result == "This sentence has bad grammar.":
        init_result["英文语法纠错模块"] = "初始化成功"
    else:
        init_result["英文语法纠错模块"] = "初始化失败"
    print("egec",init_result["英文语法纠错模块"])
    # #     英文有答案问题生成
    # eaqg_result = eaqg_process()
    # if eaqg_result == "Who starred as Etta James in Cadillac Records?":
    #     init_result["英文有答案问题生成模块"] = "初始化成功"
    # else:
    #     init_result["英文有答案问题生成模块"] = "初始化失败"
    # print("eaqg",init_result["英文有答案问题生成模块"])
    # #     中文有答案问题生成
    # caqg_result = caqg_process()
    # if caqg_result == "华沙的犹太人口占城市人口的百分之几?":
    #     init_result["中文有答案问题生成模块"] = "初始化成功"
    # else:
    #     init_result["中文有答案问题生成模块"] = "初始化失败"
    # print("caqg",init_result["中文有答案问题生成模块"])
    #     英文无答案问题生成
    enaqg_result = enaqg_process()
    if enaqg_result == "What is Hugging Face's price point?":
        init_result["英文无答案问题生成模块"] = "初始化成功"
    else:
        init_result["英文无答案问题生成模块"] = "初始化失败"
    print("enaqg",init_result["英文无答案问题生成模块"])
    #     中文无答案问题生成
    cnaqg_result = cnaqg_process()
    if cnaqg_result == "农夫在路边发现了什么?":
        init_result["中文无答案问题生成模块"] = "初始化成功"
    else:
        init_result["中文无答案问题生成模块"] = "初始化失败"
    print("cnaqg",init_result["中文无答案问题生成模块"])
    # 英文干扰项生成
    edg_result = edg_process()
    if edg_result == "traveling in France .":
        init_result["英文干扰项生成模块"] = "初始化成功"
    else:
        init_result["英文干扰项生成模块"] = "初始化失败"
    print("edg",init_result["英文干扰项生成模块"])
    # 中文干扰项生成
    cdg_result = cdg_process()
    if cdg_result == "人脸$人脸识别$图像识别":
        init_result["中文干扰项生成模块"] = "初始化成功"
    else:
        init_result["中文干扰项生成模块"] = "初始化失败"
    print("cdg",init_result["中文干扰项生成模块"])
    # 英文问答
    eqa_result = eqa_process()
    if len(str(eqa_result)) > 0:
        init_result["英文QA评估可回答性模块"] = "初始化成功"
    else:
        init_result["英文QA评估可回答性模块"] = "初始化失败"
    print("eqa",init_result["英文QA评估可回答性模块"])
    # 中文问答
    cqa_result = cqa_process()
    if str(cqa_result) == "('普希金', 0.9736020565032959)":
        init_result["中文QA评估可回答性模块"] = "初始化成功"
    else:
        init_result["中文QA评估可回答性模块"] = "初始化失败"
    print("cqa",init_result["中文QA评估可回答性模块"])
    # 一致性相关性评估
    eval_result = ceval_process()
    if eval_result is not None:
        init_result["一致性相关性评估模块"] = "初始化成功"
    else:
        init_result["一致性相关性评估模块"] = "初始化失败"
    print("eval",init_result["一致性相关性评估模块"])
    # 命名实体识别模块
    ner_result = cner_process()
    if ner_result is not None:
        init_result["命名实体识别模块"] = "初始化成功"
    else:
        init_result["命名实体识别模块"] = "初始化失败"
    print("ner",init_result["命名实体识别模块"])
    return init_result

# 模型自测
init_result = model_self_test()